#define _WIN32_DCOM
#include "pheventconsumer.h"
#include <cassert>
#include "..\..\phshared\phshared.h"
#include <comdef.h>
using namespace std;

void PHCOMError(HRESULT hr,std::string s,int line, std::string file)
{	
	std::ostringstream os;
	os<<"0x"<<std::hex<<hr;
	os<<" "<<s;
	PHTrace(os.str(),line,file);
}

// Connect to WMI 
BOOL CPHEventConsumer::Connect()
	{

	HRESULT hres;
/*CoInitializeSecurity(NULL, -1, NULL, NULL, RPC_C_AUTHN_LEVEL_NONE,
        RPC_C_IMP_LEVEL_IDENTIFY, NULL, EOAC_NONE, 0);*/
	// Release any existing connection
	/*if ( m_pSvc )
	{
		m_pSvc->Release();*/
		//m_pSvc = NULL;
	//}
	/* hres =  CoInitializeSecurity(
        NULL, 
        -1,                          // COM negotiates service
        NULL,                        // Authentication services
        NULL,                        // Reserved
        RPC_C_AUTHN_LEVEL_DEFAULT,   // Default authentication 
        RPC_C_IMP_LEVEL_IMPERSONATE, // Default Impersonation  
        NULL,                        // Authentication info
        EOAC_NONE,                   // Additional capabilities 
        NULL                         // Reserved
        );
	 if (FAILED(hres))
    {
		PHCOMError(hres,"Failed to initialize security.",__LINE__,__FILE__);
		CoUninitialize();
		return false;
	 }*/

	 IWbemLocator *pLoc;

     hres = CoCreateInstance(CLSID_WbemLocator, 0, CLSCTX_INPROC_SERVER,
            IID_IWbemLocator, (LPVOID *) &pLoc);
	 if (hres != S_OK)
    {

	    PHCOMError(hres,"CoCreateInstance on IWbemLocator failed",__LINE__,__FILE__);
		CoUninitialize();
        return FALSE;
    }	
	// Connect to the server
	hres = pLoc->ConnectServer(
             _bstr_t(L"root\\CIMV2"),//CComBSTR("\\\\.\\ROOT\\CIMV2"),
            NULL,
            NULL,
            NULL,
            0,
            0,
            0,
            &m_pSvc
            );

	pLoc->Release();

    if (FAILED(hres))
	{
  	    PHCOMError(hres,"ConnectServer failed",__LINE__,__FILE__);
		pLoc->Release();     
        CoUninitialize();
		return FALSE;
	}
	 hres = CoSetProxyBlanket(
        m_pSvc,                        // Indicates the proxy to set
        RPC_C_AUTHN_WINNT,           // RPC_C_AUTHN_xxx 
        RPC_C_AUTHZ_NONE,            // RPC_C_AUTHZ_xxx 
        NULL,                        // Server principal name 
        RPC_C_AUTHN_LEVEL_CALL,      // RPC_C_AUTHN_LEVEL_xxx 
        RPC_C_IMP_LEVEL_IMPERSONATE, // RPC_C_IMP_LEVEL_xxx
        NULL,                        // client identity
        EOAC_NONE                    // proxy capabilities 
    );
	 if (FAILED(hres))
    {
        PHCOMError(hres, "Could not set proxy blanket.",__LINE__,__FILE__);
        m_pSvc->Release();
        pLoc->Release();     
        CoUninitialize();
        return FALSE;               // Program has failed.
    }

	return TRUE;
	}
/*Create two event sinks for processes starting and stopping*/
bool CPHEventConsumer::Setup()
	{
	HRESULT hres;
	/*	if ( m_pStartSink )
	{
		// This query already running
		// cancel the query
		m_pSvc->CancelAsyncCall(m_pStartSink);
		m_pStartSink->Release();
		m_pStartSink = NULL;
		
	}*/

	
		// Checkbox is checked if we get here
		//CComBSTR Language(L"WQL");
		/*CCom*/BSTR Query;

		Query = L"SELECT * FROM Win32_ProcessStartTrace";

		 // Use an unsecured apartment for security
		IUnsecuredApartment* pUnsecApp = NULL;
		 hres = CoCreateInstance(CLSID_UnsecuredApartment, NULL, CLSCTX_LOCAL_SERVER, IID_IUnsecuredApartment, (void**)&pUnsecApp);
 
		m_pStartSink = new CALEventSink();
		m_pStartSink->AddRef();

		  IUnknown* pStubUnk = NULL; 
    pUnsecApp->CreateObjectStub(m_pStartSink, &pStubUnk);

    IWbemObjectSink* pStubSink = NULL;
    pStubUnk->QueryInterface(IID_IWbemObjectSink, (void **) &pStubSink);

		// Query the event
		/*HRESULT hrRes; = m_pSvc->ExecNotificationQueryAsync(
			Language,
			Query,
			0,                  // Flags
			0,                  // Context
			m_pStartSink
        );

		if (FAILED(hrRes))
		{
			// uh oh
			//assert(false);
			PHCOMError(hrRes,"ExecNotificationQuery failed.",__LINE__,__FILE__);
			m_pStartSink->Release();
			m_pStartSink = NULL;
			return false;
		}*/
		 hres = m_pSvc->ExecNotificationQueryAsync(
        _bstr_t("WQL"), _bstr_t(Query), WBEM_FLAG_SEND_STATUS, NULL, pStubSink);

    // Check for errors.
    if (FAILED(hres))
    {
       // printf("ExecNotificationQueryAsync failed with = 0x%X\n", hres);
		PHCOMError(hres,"ExecNotificationQuery failed.",__LINE__,__FILE__);
        m_pSvc->Release();
        //TODO:pLoc->Release();
        
        m_pStartSink->Release();
        
        CoUninitialize();    
        return false;
    }

	//pUnsecApp->Release();
        pStubUnk->Release();
		pStubSink->Release();
	/*if ( m_pStopSink )
	{
		// This query already running
		// cancel the query
		m_pSvc->CancelAsyncCall(m_pStopSink);
		m_pStopSink->Release();
		m_pStopSink = NULL;
		
	}*/

	
		// Checkbox is checked if we get here
		//CComBSTR Language(L"WQL");
		//CComBSTR Query;

		Query = L"SELECT * FROM Win32_ProcessStopTrace";
		m_pStopSink = new CALEventSink();
		m_pStopSink->AddRef();

		 // IUnknown* 
		pStubUnk = NULL; 
    pUnsecApp->CreateObjectStub(m_pStopSink, &pStubUnk);

    //IWbemObjectSink* 
		pStubSink = NULL;
    pStubUnk->QueryInterface(IID_IWbemObjectSink, (void **) &pStubSink);
		// Query the event
		/*HRESULT hrRes = m_pSvc->ExecNotificationQueryAsync(
			Language,
			Query,
			0,                  // Flags
			0,                  // Context
			m_pStopSink
        );

		if (FAILED(hrRes))
		{
			// uh oh
			//assert(false);
			PHCOMError(hrRes,"ExecNotificationQuery failed.",__LINE__,__FILE__);
			m_pStopSink->Release();
			m_pStopSink = NULL;
			return false;
		}		*/
	 hres = m_pSvc->ExecNotificationQueryAsync(
        _bstr_t("WQL"), _bstr_t(Query), WBEM_FLAG_SEND_STATUS, NULL, pStubSink);

    // Check for errors.
    if (FAILED(hres))
    {
       // printf("ExecNotificationQueryAsync failed with = 0x%X\n", hres);
		PHCOMError(hres,"ExecNotificationQuery failed.",__LINE__,__FILE__);
        m_pSvc->Release();
        //todo:pLoc->Release();
        //pUnsecApp->Release();
        //pStubUnk->Release();
        m_pStopSink->Release();
        //pStubSink->Release();
        CoUninitialize();    
        return false;
    }
	pUnsecApp->Release();
        pStubUnk->Release();
		pStubSink->Release();
		return true;
	}